//==========================================================================
// Copyright (c) 2000-2008,  Elastos, Inc.  All Rights Reserved.
//==========================================================================
#include <winhack.h>
//#include <elasys_server.h>
#include <sysconfig.h>
#include <CMalloc.h>
//#include <pal.h>

#define PAGESIZE 0x1000
#define PAGESIZE_ROUNDOFF 0xfff
#define RoundPages(size) ((size + PAGESIZE_ROUNDOFF) & (~PAGESIZE_ROUNDOFF))
#define HasRoundPages(Size) ((Size & 0xfff) == 0)
#define HasRoundAddress(Address) (((DWORD)Address & 0xfff) == 0)

//
//
// CPreBChkMalloc:
// (Previous Memory Bounds Checker Malloc Object)
// When reading or writing before the heap block,
// The system will catch the error right now!
//
// Theory:
// For each heap block, commit 1 or n physical pages,
// and the block occupies the part from the beginning.
// And the previous one page is reserved.
// If reading or writing before the heap block,
// Probably the address belongs to the previous one page.
// There will be a error detected right now,
// since the page is only reserved, not commited.
//

// | 1 Page (Reserved, NOT Commited) |        1 ~ n Pages (Commited)       |
// +---------------------------------+----------------+--------------------+
// |              RESERVED           | User Content   |     Not Used       |
// +---------------------------------+----------------+--------------------+
// |                                 ^                                     |
//                                   |
//                                   |__Return Address For Users

PVoid CPreBChkMalloc::Alloc(size_t cb)
{
    byte *pAllocMem;
    size_t TotalSize;

    TotalSize = RoundPages(cb);

    pAllocMem =  (byte *)VirtualAlloc(NULL, TotalSize + PAGESIZE,
                                 MEM_RESERVE, PAGE_NOACCESS);
    if (NULL == pAllocMem) return NULL;

    pAllocMem = (byte *)VirtualAlloc(pAllocMem + PAGESIZE, TotalSize,
                                 MEM_COMMIT, PAGE_READWRITE);
    if (NULL == pAllocMem) {
        VirtualFree(pAllocMem, 0 ,MEM_RELEASE);
        return NULL;
    }
    return (PVoid)pAllocMem;
}

PVoid CPreBChkMalloc::Realloc(PVoid pv, size_t cb)
{
    if (pv) {
        assert(HasRoundAddress(pv));
        if (pv && cb == 0) {
            CPreBChkMalloc::Free(pv);
            return NULL;
        }
        VirtualFree((byte *)pv - PAGESIZE, 0, MEM_DECOMMIT | MEM_RELEASE);
    }

    return CPreBChkMalloc::Alloc(cb);
}

void CPreBChkMalloc::Free(PVoid pv)
{
    if (NULL == pv) {
        //wprintf(L"try to free a NULL pointer!\n");
        return;
    }

    assert(HasRoundAddress(pv));
    VirtualFree((byte *)pv, 1 , MEM_DECOMMIT);
    VirtualFree((byte *)pv - PAGESIZE, 0, MEM_RELEASE);
}

//
//
// CPostBChkMalloc:
// (Post Memory Bounds Checker Malloc Object)
// When reading or writing after the heap block,
// The system will catch the error right now!
//
// Theory:
// For each heap block, commit 1 or n physical pages,
// and the block occupies the tail of these pages.
// And the next one page is reserved.
// If reading or writing after the heap block,
// Probably the address belongs to the next one page.
// There will be a error detected right now,
// since the page is only reserved, not commited.
//

// |        1 ~ n Pages (Commited)       | 1 Page (Reserved, NOT Commited) |
// +----------------+------+-------------+---------------------------------+
// | Not Used       |Header| User Content|              RESERVED           |
// +----------------+------+-------------+---------------------------------+
// |                       ^             |
//                         |
//                         |__Return Address For Users

// NOTE: It needs lots of memory! So it should be only used when necessary.

typedef struct _BChkHeader {
    PVoid pAllocPage; // Pointer to the address of 1st page.
    DWORD Size;       // The size of this heap block,
                      // 8 bytes rounded,  including header.
} BChkHeader, *PBChkHeader;

PVoid CPostBChkMalloc::Alloc(size_t cb)
{
    PVoid pRetMem;
    PBChkHeader pChkHeader;
    byte *pAllocMem;
    size_t TotalSize;

    TotalSize = RoundPages(cb + sizeof(BChkHeader));
    pAllocMem =  (byte *)VirtualAlloc(NULL, TotalSize + PAGESIZE,
                                MEM_RESERVE, PAGE_NOACCESS);
    if (NULL == pAllocMem) return NULL;

    pAllocMem = (byte *)VirtualAlloc(pAllocMem, TotalSize,
                                MEM_COMMIT, PAGE_READWRITE);
    if (NULL == pAllocMem) {
        VirtualFree(pAllocMem, 0, MEM_DECOMMIT);
        return NULL;
    }

    pRetMem = (PVoid)(pAllocMem + TotalSize - cb);
    pChkHeader = (PBChkHeader)pRetMem - 1;
    pChkHeader->pAllocPage = pAllocMem;
    pChkHeader->Size = cb + sizeof(BChkHeader);

    return pRetMem;
}

PVoid CPostBChkMalloc::Realloc(PVoid pv, size_t cb)
{
    PVoid pOldAllocPage = NULL;
    PVoid pRetMem;
    size_t OldUserSize = 0, NewUserSize = 0;

    if (pv) {
        PBChkHeader pOldHeader;
        size_t OldTotalSize, NewTotalSize;

        if (pv && cb == 0) {
            CPostBChkMalloc::Free(pv);
            return NULL;
        }

        pOldHeader = (PBChkHeader)pv - 1;
        OldUserSize = pOldHeader->Size - sizeof(BChkHeader);
        OldTotalSize = RoundPages(pOldHeader->Size);
        pOldAllocPage = pOldHeader->pAllocPage;

        NewUserSize = cb;

        if (NewUserSize == OldUserSize) {
            return pv;
        }

        NewTotalSize = RoundPages(NewUserSize + sizeof(BChkHeader));

        if (NewTotalSize <= OldTotalSize) {
            size_t FreeSize = OldTotalSize - NewTotalSize;
            PVoid pNewMem;
            PBChkHeader pNewHeader;

            assert(HasRoundAddress((byte *)pv + OldUserSize));
            pNewMem = (byte *)pv + OldUserSize - NewUserSize;
            pNewHeader = (PBChkHeader)pNewMem - 1;
            pNewHeader->Size = NewUserSize + sizeof(BChkHeader);
            pNewHeader->pAllocPage = (byte *)pOldAllocPage + FreeSize;

            memmove(pNewMem, pv, NewUserSize > OldUserSize ? OldUserSize: NewUserSize);

            if (FreeSize > 0) {
                assert(HasRoundPages(FreeSize));
                VirtualFree(pOldAllocPage, FreeSize, MEM_DECOMMIT);
                VirtualFree(pOldAllocPage, 0, MEM_DECOMMIT);
            }
            return pNewMem;
        }
    }

    pRetMem = CPostBChkMalloc::Alloc(cb);

    if (pRetMem && pOldAllocPage) {
        assert(OldUserSize > 0 && NewUserSize >  0);
        assert(OldUserSize < NewUserSize);
        memcpy(pRetMem, pv, OldUserSize);
        VirtualFree(pOldAllocPage, 0, MEM_DECOMMIT | MEM_DECOMMIT);
    }

    return pRetMem;
}

void CPostBChkMalloc::Free(PVoid pv)
{
    PVoid pAllocMem;
    PBChkHeader pHeader;

    if (NULL == pv) {
        //wprintf(L"try to free a NULL pointer!\n");
        return;
    }

    pHeader = (PBChkHeader)pv - 1;
    pAllocMem = pHeader->pAllocPage;

    assert(HasRoundAddress(pAllocMem));
    VirtualFree(pv, pHeader->Size + sizeof(BChkHeader), MEM_DECOMMIT);
    VirtualFree(pAllocMem, 0, MEM_DECOMMIT);
}
